#!/usr/bin/env python3
"""Auto Video - 文字生成图片主程序

命令行接口，支持通过参数生成图片。
"""

import os
import sys
import time
import hashlib
from datetime import datetime
from pathlib import Path
from typing import Optional

import click
from PIL import Image

# 添加src目录到Python路径
sys.path.insert(0, str(Path(__file__).parent / "src"))

from src.generators.text_to_image import TextToImage
from src.utils.config import get_config


def generate_filename(prompt: str, template: str = "{timestamp}_{prompt_hash}", ext: str = "png") -> str:
    """生成文件名
    
    Args:
        prompt: 提示词
        template: 文件名模板
        ext: 文件扩展名
        
    Returns:
        生成的文件名
    """
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    prompt_hash = hashlib.md5(prompt.encode('utf-8')).hexdigest()[:8]
    
    filename = template.format(
        timestamp=timestamp,
        prompt_hash=prompt_hash,
        prompt=prompt[:20].replace(" ", "_")  # 截取前20个字符作为文件名
    )
    
    return f"{filename}.{ext}"


@click.command()
@click.option(
    "--prompt", "-p",
    required=True,
    help="图片生成的文字描述"
)
@click.option(
    "--output", "-o",
    help="输出文件路径（可选，默认使用配置中的输出目录）"
)
@click.option(
    "--negative-prompt", "-n",
    help="负向提示词（可选）"
)
@click.option(
    "--width", "-w",
    type=int,
    help="图片宽度（默认512）"
)
@click.option(
    "--height", "-h",
    type=int,
    help="图片高度（默认512）"
)
@click.option(
    "--steps", "-s",
    type=int,
    help="推理步数（默认50）"
)
@click.option(
    "--guidance-scale", "-g",
    type=float,
    help="引导强度（默认7.5）"
)
@click.option(
    "--seed",
    type=int,
    help="随机种子（可选）"
)
@click.option(
    "--num-images", "-c",
    type=int,
    default=1,
    help="生成图片数量（默认1）"
)
@click.option(
    "--model", "-m",
    help="模型ID或路径（可选）"
)
@click.option(
    "--device",
    type=click.Choice(['auto', 'cuda', 'cpu', 'mps']),
    help="计算设备（可选）"
)
@click.option(
    "--config", "-cfg",
    help="配置文件路径（可选）"
)
@click.option(
    "--verbose", "-v",
    is_flag=True,
    help="详细输出"
)
def main(
    prompt: str,
    output: Optional[str],
    negative_prompt: Optional[str],
    width: Optional[int],
    height: Optional[int],
    steps: Optional[int],
    guidance_scale: Optional[float],
    seed: Optional[int],
    num_images: int,
    model: Optional[str],
    device: Optional[str],
    config: Optional[str],
    verbose: bool
):
    """Auto Video - 文字生成图片工具
    
    使用 Stable Diffusion 将文字描述转换为图片。
    
    示例:
        python main.py -p "一只可爱的小猫在花园里玩耍"
        python main.py -p "sunset over mountains" -o "./my_image.png" -w 768 -h 768
    """
    try:
        # 加载配置
        if verbose:
            print("正在加载配置...")
        
        cfg = get_config(config)
        
        # 获取配置参数
        model_config = cfg.get_model_config()
        generation_config = cfg.get_generation_config()
        output_config = cfg.get_output_config()
        
        # 命令行参数覆盖配置文件
        model_id = model or model_config.get('name')
        device_name = device or model_config.get('device')
        
        generation_params = {
            'width': width or generation_config.get('width', 512),
            'height': height or generation_config.get('height', 512),
            'num_inference_steps': steps or generation_config.get('num_inference_steps', 50),
            'guidance_scale': guidance_scale or generation_config.get('guidance_scale', 7.5),
            'seed': seed or generation_config.get('seed'),
            'num_images': num_images,
            'negative_prompt': negative_prompt or cfg.get('prompts.negative_prompt')
        }
        
        if verbose:
            print(f"模型: {model_id}")
            print(f"设备: {device_name}")
            print(f"生成参数: {generation_params}")
        
        # 初始化生成器
        print("正在初始化文字转图片生成器...")
        generator = TextToImage(
            model_id=model_id,
            device=device_name if device_name != 'auto' else None
        )
        
        # 生成图片
        print(f"\n开始生成图片...")
        print(f"提示词: {prompt}")
        
        start_time = time.time()
        
        images = generator.generate(
            prompt=prompt,
            **generation_params
        )
        
        end_time = time.time()
        generation_time = end_time - start_time
        
        # 处理输出
        if not isinstance(images, list):
            images = [images]
        
        saved_files = []
        
        for i, image in enumerate(images):
            # 确定输出路径
            if output:
                if num_images > 1:
                    # 多张图片时添加序号
                    output_path = Path(output)
                    output_file = output_path.parent / f"{output_path.stem}_{i+1}{output_path.suffix}"
                else:
                    output_file = Path(output)
            else:
                # 使用默认输出目录
                output_dir = Path(output_config.get('directory', './outputs'))
                output_dir.mkdir(parents=True, exist_ok=True)
                
                filename_template = output_config.get('filename_template', '{timestamp}_{prompt_hash}')
                file_format = output_config.get('format', 'png')
                
                if num_images > 1:
                    filename = generate_filename(f"{prompt}_{i+1}", filename_template, file_format)
                else:
                    filename = generate_filename(prompt, filename_template, file_format)
                
                output_file = output_dir / filename
            
            # 保存图片
            generator.save_image(
                image,
                str(output_file),
                quality=output_config.get('quality', 95)
            )
            
            saved_files.append(str(output_file))
        
        # 输出结果
        print(f"\n✅ 生成完成！")
        print(f"⏱️  用时: {generation_time:.2f} 秒")
        print(f"📁 保存位置:")
        for file_path in saved_files:
            print(f"   {file_path}")
        
        # 清理GPU缓存
        generator.clear_cache()
        
    except KeyboardInterrupt:
        print("\n❌ 用户中断操作")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 生成失败: {e}")
        if verbose:
            import traceback
            traceback.print_exc()
        sys.exit(1)


@click.group()
def cli():
    """Auto Video 命令行工具集"""
    pass


@cli.command()
@click.option("--config", "-c", help="配置文件路径")
def info(config: Optional[str]):
    """显示系统信息"""
    import torch
    
    cfg = get_config(config)
    
    print("🖥️  系统信息:")
    print(f"   Python: {sys.version.split()[0]}")
    print(f"   PyTorch: {torch.__version__}")
    print(f"   CUDA可用: {'是' if torch.cuda.is_available() else '否'}")
    
    if torch.cuda.is_available():
        print(f"   GPU数量: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            gpu_name = torch.cuda.get_device_name(i)
            gpu_memory = torch.cuda.get_device_properties(i).total_memory / 1024**3
            print(f"   GPU {i}: {gpu_name} ({gpu_memory:.1f}GB)")
    
    print(f"\n⚙️  当前配置:")
    model_config = cfg.get_model_config()
    generation_config = cfg.get_generation_config()
    
    print(f"   模型: {model_config.get('name')}")
    print(f"   设备: {model_config.get('device')}")
    print(f"   分辨率: {generation_config.get('width')}x{generation_config.get('height')}")
    print(f"   推理步数: {generation_config.get('num_inference_steps')}")


@cli.command()
@click.argument('prompt')
@click.option('--output', '-o', help='输出路径')
@click.option('--config', '-c', help='配置文件路径')
def generate(prompt: str, output: Optional[str], config: Optional[str]):
    """快速生成图片（简化版命令）"""
    # 调用主生成函数
    ctx = click.Context(main)
    ctx.invoke(main, prompt=prompt, output=output, config=config, 
              negative_prompt=None, width=None, height=None, steps=None,
              guidance_scale=None, seed=None, num_images=1, model=None,
              device=None, verbose=False)


if __name__ == "__main__":
    # 如果直接运行main.py，使用命令组
    if len(sys.argv) == 1:
        # 没有参数时显示帮助
        cli(['--help'])
    else:
        cli()